import gradio as gr
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans

def extract_colors(image, num_colors=5):
    # 将图片转换为RGB模式并调整大小
    image = image.convert('RGB')
    image = image.resize((100, 100))  # 调整大小以加快处理速度
    pixels = np.array(image).reshape(-1, 3)  # 将图像数据转换为二维数组

    # 使用K-Means聚类提取主色调
    kmeans = KMeans(n_clusters=num_colors)
    kmeans.fit(pixels)
    colors = kmeans.cluster_centers_.astype(int)  # 获取主色调
    return colors

def colors_to_hex(colors):
    # 将颜色元组转换为十六进制表示
    return [f"#{r:02x}{g:02x}{b:02x}" for r, g, b in colors]

def analyze_image(image):
    # 提取颜色
    colors = extract_colors(image)
    # 转换颜色为十六进制形式
    hex_colors = colors_to_hex(colors)
    
    # 创建调色板图像
    palette = Image.new("RGB", (100, 30))  # 创建一个100x50的空白图像
    for i, color in enumerate(colors):
        palette.paste(Image.new("RGB", (20, 30), tuple(color)), (i * 20, 0))  # 每个颜色占20x50的区域
    
    # 返回颜色及其十六进制表示和调色板图像
    return "\n".join(hex_colors), palette

# 创建Gradio界面
iface = gr.Interface(
    fn=analyze_image, 
    inputs=gr.inputs.Image(type="pil"), 
    outputs=[
        gr.outputs.Textbox(label="色彩十六进制码"), 
        gr.outputs.Image(label="Color Palette")
    ],
    title="Kmeans调色板提取器"
)

# 启动应用
iface.launch()